import base64
import hmac
import time
import json


class MyJwt:
    """
    实例化JWT算法
    """

    def bs64(self, tem):
        """生成header bs"""
        header_str = json.dumps(tem)
        header_bs = base64.b64encode(header_str.encode())
        return header_bs

    def _encode(self, algorithm, key_, _payload):
        """
        加密生成token
        :param algorithm: 加密方式
        :param key: 公司密钥
        :param payload: 私钥
        :return: token
        """
        header = {"alg": algorithm, "typ": "JWT"}
        header_ = self.bs64(header)
        payload_ = self.bs64(_payload)
        msg = header_ + b"." + payload_
        sign_str = hmac.new(key=key_.encode(), msg=msg, digestmod="SHA256").hexdigest()
        sign_bs = base64.b64encode(sign_str.encode())
        token = msg + b"." + sign_bs
        return token

    def _decode(self, token, key):
        """
        对token加密校验
        :param token: 要校验的token
        :param key: 公司密钥
        :param algorithms: 加密方式
        :return:
        """
        header, _payload, sign = token.split(b".")
        sign_user = hmac.new(key=key.encode(), msg=header + b"." + _payload, digestmod="SHA256").hexdigest()
        sign_user = base64.b64encode(sign_user.encode())
        if sign == sign_user:
            return True
        else:
            return False


if __name__ == '__main__':
    payload = {
        "exp": int(time.time()) + 5,
        "username": "zhaoliying",
    }
    key = "123456"
    my = MyJwt()

    token = my._encode(algorithm="HS256", _payload=payload, key_=key)
    print(token)

    if my._decode(token=token, key=key):
        print("success")
    else:
        print('default')
